#include <iostream>

#include "arenas/arena_base.h"
#include "output_base.h"
#include "machine.h"
#include "predictors/predictor_target.h"
#include "predictors/predictor_const_v.h"
#include "model.h"
#include "predictors/predictor_link.h"
#include "ioconnections.h"
#include "trace.h"

using namespace std;

Machine::Machine()
{
  _previous_time = 0.0;
  _arena = NULL;
  _trace = NULL;
}

Machine::~Machine()
{
}

void Machine::init(void)
{
  if (!_arena) {
    cout << "ERROR: _arena is NULL" << endl;
    exit(1);
  }
  // Create arena and set arena state + machine initial state
  update_input(0.0);
  save_state(0.0);
}

void Machine::reset(void)
{
//  cout << "MACHINE: RESET" << endl;
  // reset predictors
  for (int i = 0; i < (int) _models.size(); ++i) {
    _models[i]->reset();
  }
  // reset arena
  _arena->reset();
  _arena->update(0.0); // set arena state to time = 0.0
  // reset time
  _previous_time = 0.0;
  // reset input links
  // TODO: reset links
}

OConnection *Machine::create_output_connection(Model *m, boost::function<void(
    float)> f)
{
  OConnection *oc = new OConnection(this, m);
  oc->set_function(f);

  return oc;
}

OConnection *Machine::create_output_connection(Model *m, boost::function<void(
    float)> f, string name, Output_base *output_module)
{
  OConnection *oc = create_output_connection(m, f);
  oc->set_name(name);
  oc->set_output(output_module);

  return oc;
}

IConnection *Machine::create_input_connection(Predictor_base *p,
    int input_index, float *source)
{
  IConnection *ic = new IConnection(p->get_input_link(input_index));
  ic->set_value_p(source);
  add_input_connection(ic);

  return ic;
}

Model *Machine::get_model(int i)
{
  if (i < (int) _models.size()) {
    return _models[i];
  }
  cout << "ERROR: Index out of range: " << i << endl;
  exit(1);
}

std::vector<Model *> *Machine::get_models(void)
{
  return &_models;
}

void Machine::add_model(Model *m)
{
  _models.push_back(m);
}

void Machine::add_output_connection(OConnection *oc)
{
  _output_connections.push_back(oc);
}

void Machine::add_input_connection(IConnection *ic)
{
  _input_connections.push_back(ic);
}

Arena_base *Machine::get_arena(void)
{
  return _arena;
}

void Machine::set_arena(Arena_base *arena)
{
  _arena = arena;
}

void Machine::predict(float t)
{
//  cout << "PREDICT " << t << endl;
  for (int i = 0; i < (int) _models.size(); ++i) {
    _models[i]->predict(t);
  }
}

void Machine::activate_models(float t)
{
//  cout << "ACTIVATE_MODELS " << endl;
  for (int i = 0; i < (int) _models.size(); ++i) {
    _models[i]->calculate_error(t);
    _models[i]->activate(0.5); // argument: critical error
  }
}

// Activate output links if predictor is active.
// This step must be separate from activate_predictors() since
// activate_predictors() depends on output links.
void Machine::activate_model_output_links(void)
{
//  cout << "ACTIVATE_PREDICTORS_OUTPUT_LINKS " << endl;
  for (int i = 0; i < (int) _models.size(); ++i) {
    _models[i]->activate_output_link();
  }
}

void Machine::save_state(float t)
{
  for (int i = 0; i < (int) _models.size(); ++i) {
    _models[i]->save_state(t);
  }
}

void Machine::update_arena(float t)
{
//  cout << "UPDATE_ARENA " << endl;
  _arena->update(t);
}

void Machine::update_input(float t)
{
  cout << "UPDATE_INPUT " << t << endl;
  for (int i = 0; i < (int) _input_connections.size(); ++i) {
    _input_connections[i]->get_value();
    cout << *(_input_connections[i]->get_value_p()) << endl;
  }
}

vector<OConnection *> &Machine::get_output_connections()
{
 return _output_connections;
}

vector<IConnection *> &Machine::get_input_connections()
{
 return _input_connections;
}

void Machine::prepare_arena(void)
{
//  cout << "PREPARE_ARENA " << endl;
  for (int i = 0; i < (int) _output_connections.size(); ++i) {
    _output_connections[i]->set_value();
  }
  for (int i = 0; i < (int) _input_connections.size(); ++i) {
    _input_connections[i]->get_value();
  }
}

void Machine::set_trace(Trace *trace)
{
  _trace = trace;
}

void Machine::remove_trace(void)
{
  _trace = NULL;
}

void Machine::update_trace(double t)
{
  if(_trace == NULL) {
    cout << "ERROR: trace must be set before calling update_trace()" << endl;
    exit(1);
  }
  _trace->update(t);
}

void Machine::step(float t)
{
//  cout << "STEP from " << _previous_time << " to " << t << endl;
  if (t <= _previous_time) {
    cout << "ERROR: Time must increase: " << t << " <= " << _previous_time
        << endl;
    exit(1);
  }
  // predict input to all predictors at time t
  predict(t);

  // set arena to time t
  update_arena(t);

  // set input to input at time t
  update_input(t);

  if(_trace != NULL)
    update_trace(t);

  // activate predictors
  activate_models(t);

  // implement active predictors
  activate_model_output_links();

  // save state
  save_state(t);

  // set arena state variables
  prepare_arena();

  // save time of last update
  _previous_time = t;
}
